#ifndef CODEGENERATOR_H
#define CODEGENERATOR_H

#include <string>
#include <vector>
#include <map>
#include <stack>
#include <memory>
#include <unordered_map>
#include "ASTNode.h"
#include "SymbolTable.h"

// RISC-V寄存器定义
enum class Register {
    ZERO,   // x0 - 零寄存器
    RA,     // x1 - 返回地址
    SP,     // x2 - 栈指针
    GP,     // x3 - 全局指针
    TP,     // x4 - 线程指针
    T0,     // x5 - 临时寄存器
    T1,     // x6 - 临时寄存器
    T2,     // x7 - 临时寄存器
    S0,     // x8 - 保存寄存器
    S1,     // x9 - 保存寄存器
    A0,     // x10 - 参数/返回值寄存器
    A1,     // x11 - 参数寄存器
    A2,     // x12 - 参数寄存器
    A3,     // x13 - 参数寄存器
    A4,     // x14 - 参数寄存器
    A5,     // x15 - 参数寄存器
    A6,     // x16 - 参数寄存器
    A7,     // x17 - 参数寄存器
    S2,     // x18 - 保存寄存器
    S3,     // x19 - 保存寄存器
    S4,     // x20 - 保存寄存器
    S5,     // x21 - 保存寄存器
    S6,     // x22 - 保存寄存器
    S7,     // x23 - 保存寄存器
    S8,     // x24 - 保存寄存器
    S9,     // x25 - 保存寄存器
    S10,    // x26 - 保存寄存器
    S11,    // x27 - 保存寄存器
    T3,     // x28 - 临时寄存器
    T4,     // x29 - 临时寄存器
    T5,     // x30 - 临时寄存器
    T6      // x31 - 临时寄存器
};

// 寄存器分配器
class RegisterAllocator {
private:
    std::map<Register, bool> registerUsage;
    std::map<std::string, Register> variableRegisters;
    std::stack<Register> freeRegisters;
    
public:
    RegisterAllocator();
    Register allocateRegister();
    void freeRegister(Register reg);
    Register getVariableRegister(const std::string& varName);
    void setVariableRegister(const std::string& varName, Register reg);
    void clearVariableRegisters();
    std::string registerToString(Register reg);
};

// 代码生成器
class CodeGenerator {
private:
    std::vector<std::string> code;
    RegisterAllocator regAlloc;
    int labelCounter;
    int stackOffset;
    std::string currentFunction;
    int currentFrameSize;
    std::map<std::string, int> functionStackSizes;
    std::stack<std::string> breakLabels; // break标签栈
    std::stack<std::string> continueLabels; // continue标签栈
    std::unordered_map<std::string, int> varOffsets; // 变量偏移表
    int nextOffset = -4; // 当前函数变量偏移
    
    // 辅助方法
    std::string generateLabel(const std::string& prefix);
    void emit(const std::string& instruction);
    void emitComment(const std::string& comment);
    void emitFunctionPrologue(const std::string& funcName, int frameSize);
    void emitFunctionEpilogue(int frameSize);
    void saveRegisters();
    void restoreRegisters();
    
    // 表达式代码生成
    Register generateExpression(ASTNode* node);
    Register generateBinaryOp(ASTNode* node);
    Register generateUnaryOp(ASTNode* node);
    Register generateVarRef(ASTNode* node);
    Register generateNumLiteral(ASTNode* node);
    Register generateFuncCall(ASTNode* node);
    
    // 语句代码生成
    void generateStatement(ASTNode* node, bool& hasReturn, int frameSize);
    void generateAssignStmt(ASTNode* node);
    void generateDeclAssignStmt(ASTNode* node);
    void generateIfStmt(ASTNode* node, bool& hasReturn, int frameSize);
    void generateIfElseStmt(ASTNode* node, bool& hasReturn, int frameSize);
    void generateWhileStmt(ASTNode* node, bool& hasReturn, int frameSize);
    void generateBreakStmt(ASTNode* node);
    void generateContinueStmt(ASTNode* node);
    void generateReturnStmt(ASTNode* node, bool& hasReturn, int frameSize);
    void generateBlock(ASTNode* node, bool& hasReturn, int frameSize);
    
    // 条件代码生成
    void generateCondition(ASTNode* node, const std::string& trueLabel, const std::string& falseLabel);
    void generateShortCircuit(ASTNode* node, const std::string& trueLabel, const std::string& falseLabel);
    
    // 栈管理
    void pushToStack(Register reg);
    Register popFromStack();
    int getVariableOffset(const std::string& varName);
    
    int countLocalVars(ASTNode* node);
    
public:
    CodeGenerator();
    void generate(ASTNode* root);
    void printCode() const;
    void saveToFile(const std::string& filename) const;
    
    // 优化方法
    void optimizeCommonSubexpressions();
    void optimizeConstantFolding();
    void optimizeRegisterAllocation();
};

#endif // CODEGENERATOR_H 